package org.gjy.m8.bucket;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Leaky Bucket（平滑限流），适合“每秒固定速率放行”场景。
 * <p>
 * 5.1 思路（固定泄漏速率）
 * 桶的 容量 为 capacity，泄漏速率 为 rate（每秒流出多少请求）。
 * 当请求到来：
 * 先泄漏（依据时间间隔把水位降低）。
 * 若 水位 < capacity → 水位 +1，通过；否则 拒绝。
 *
 * @author gjy
 * @version 1.0
 * @since 2025-10-13 17:59:09
 */
public class LeakyBucketRateLimiter<K> implements RateLimiter<K> {

    private static class Bucket {
        final long capacity;          // 最大水位
        final double rate;            // 每秒泄漏的请求数
        final AtomicLong lastLeakTs = new AtomicLong(System.nanoTime()); // 纳秒
        final AtomicLong water = new AtomicLong(0L);                    // 当前水位（可能为小数，使用 double 乘 1000）

        Bucket(long capacity, double rate) {
            this.capacity = capacity;
            this.rate = rate;
        }

        /**
         * 计算并执行泄漏
         */
        private void leak() {
            long now = System.nanoTime();
            long prev = lastLeakTs.getAndSet(now);
            long elapsedNs = now - prev;
            double leaked = (elapsedNs / 1_000_000_000.0) * rate;   // 计算漏掉的请求数量
            if (leaked > 0) {
                // 用 CAS 把 water 减去 leaked（可能为小数，需要放大）
                long decrement = (long) (leaked * 1000L); // 放大 1000
                water.updateAndGet(v -> Math.max(0L, v - decrement));
            }
        }

        boolean tryConsume() {
            leak(); // 先把旧水位泄漏掉
            while (true) {
                long cur = water.get();
                if (cur >= capacity * 1000L) {
                    return false; // 桶已满，限流
                }
                // 尝试加 1（放大 1000）
                if (water.compareAndSet(cur, cur + 1000L)) {
                    return true; // 成功放入桶
                }
            }
        }
    }

    private final ConcurrentMap<K, Bucket> buckets = new ConcurrentHashMap<>();
    private final long capacity;
    private final double rate; // QPS（每秒允许的请求数）

    public LeakyBucketRateLimiter(long capacity, double rate) {
        if (capacity <= 0 || rate <= 0) {
            throw new IllegalArgumentException("capacity & rate 必须 > 0");
        }
        this.capacity = capacity;
        this.rate = rate;
    }

    @Override
    public boolean tryAcquire(K key) {
        Bucket b = buckets.computeIfAbsent(key,
                k -> new Bucket(capacity, rate));
        return b.tryConsume();
    }

}
